#coding=utf-8
import json
import pymongo
import copy
from bson import json_util
from Core.DataSeries import *
import Core.Gadget as Gadget
import Core.DataBase as DataBase

MAX_DATABASE_SCAN = 100000000


class MongoDB(DataBase.DataBase):

    def __init__(self, address, port, username="root", password=None):

        print("Connecting MongoDB")
        self.address = address
        self.port = port
        self.username = username
        self.password = password
        # Create a Connection "mongodb://mongodb0.example.net:27019"
        self.client = pymongo.MongoClient("mongodb://" + address + ":" + port)
        if self.password != None:
            db = self.client["admin"]
            db.authenticate(self.username, self.password, mechanism='SCRAM-SHA-1')
        print("Connected MongoDB")


    def Count(self, databaseName, collectionName, filter={}):

        # ---Old Fashion---
        # if len(filter) == 0:
        #    db = self.client[databaseName]
        #    # coll = db[collectionName].count(filter)
        #    ret = db.command("collstats", collectionName)
        #    return ret["count"]
        #else:
            # docs = self.Find(databaseName, collectionName, filter)
            # return len(docs)

        skip = 0
        limit = 0
        if "skip" in filter:
            skip = filter["skip"]
        if "limit" in filter:
            limit = filter["limit"]
        db = self.client[databaseName]
        cursor = db[collectionName].find(filter, skip=skip, limit=limit).max_scan(MAX_DATABASE_SCAN).count()
        return cursor


    def Find(self, databaseName, tableName, filter={}, sort={}, projection={}):
        #
        if tableName.lower() == "index":
            tableName = "index2"
        #
        sortDoc = None
        if len(sort) == 0:
            sortDoc = None
        else:
            sortDoc = []
            for item in sort:
                key = item[0]
                value = item[1]
                if value == 1:
                    sortMethod = pymongo.ASCENDING
                elif value == 0:
                    sortMethod = pymongo.DESCENDING
                sortDoc.append((key, sortMethod))
        #
        return self.findWithFilter(databaseName, tableName, filter, sortDoc, projection)


    def GetDataSeries(self, symbol, dataType=None, instrumentType=None, datetime1=None, datetime2=None):

        # ---Old Fashion---
        # dataSeriesName = symbol + "_Time_86400_Bar"
        # dataSeries = self.getDataSeries(dataSeriesName, datetime1, datetime2, instrumentType)
        # return dataSeries
        # if instrumentType == "Bond" or instrumentType == "Option":

        #
        filter = {}
        filter["Symbol"] = symbol
        filter["StdDateTime"] = {}
        if datetime1 != None:
            filter["StdDateTime"]["$gte"] = datetime1
        if datetime2 != None:
            filter["StdDateTime"]["$lte"] = datetime2
        if not filter["StdDateTime"]:
            del filter["StdDateTime"]

        databaseName = ""
        collectionName = ""
        dataSeriesName = ""

        sort = [("StdDateTime", 1)]
        if dataType == "DailyBar":
            databaseName = instrumentType
            collectionName = "DailyBar"
            dataSeriesName = symbol + "_Time_86400_Bar"

        elif dataType == "Factor":
            databaseName = "Factor"
            collectionName = instrumentType
            dataSeriesName = symbol + "_" + instrumentType

        # ---Load Database---
        series = self.Find(databaseName, collectionName, filter, sort)

        # ---Build DataSeries Object---
        dataSeries = DataSeries(dataSeriesName)
        for d in series:
            # print(document)
            if "_id" in d:
                del d["_id"]
            dataSeries.Add(d)

        #
        return dataSeries


    def SaveDataSeries(self, dataSeries, dataType=None, instrumentType=None, isUpSert=True):

        #
        dataSeriesName = dataSeries.model_name
        print(str(datetime.now()) + " Saving DataSeries " + dataSeriesName)
        #
        if dataSeries.Count() == 0:
            print(dataSeriesName + " saveDataSeries::No Data to Save")
            return

        symbol = DataSeries.ParseSymbol(dataSeriesName)
        name = DataSeries.ParseDataName(dataSeriesName)

        if dataType == None:
            dataType = DataSeries.ParseDataType(dataSeriesName)

        databaseName = ""
        collectionName = ""
        keys = []
        insertType = 0
        #
        if dataType == "DailyBar":
            databaseName = instrumentType
            collectionName = dataType
            self.saveCollection(dataSeries.items, databaseName, collectionName)
        #
        elif dataType == "Bar":
            size = self.parseSize(dataSeriesName)
            dataBaseStructureType = "BySymbol"
            if instrumentType == None:
                # stockList = self.find("Instruments","Stock",fields={"Symbol":1})
                print("Unknown Bar Type")
                return
            elif (instrumentType == "Stock") and (size == 86400 or size == 60 or size == 300):
                databaseName = "Quote"
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, databaseName, collectionName)
            elif instrumentType == "Index":
                databaseName = "Index"
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, databaseName, collectionName)
            elif instrumentType == "Future":
                databaseName = "Future"
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, databaseName, collectionName)
            elif instrumentType == "Bond":
                databaseName = "Bond"
                # query["DataSeries"] = dataSeriesName
            elif instrumentType == "MutualFund":
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, databaseName, collectionName)
            elif instrumentType == "StockOption":
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, "Option", "Dailybar")
            else:
                print("Unknown Bar Type")
                return

        elif dataType == "Fundamental":
            databaseName = "Fundamental"
            collectionName = dataSeriesName
            self.saveCollection(dataSeries.items, databaseName, collectionName)

        elif dataType == "InstrumentList":
            databaseName = "Instruments"
            collectionName = "InstrumentList"

        elif dataType == "Factor":
            databaseName = "Factor"
            collectionName = name
            self.saveCollection(dataSeries.items, databaseName, collectionName, isUpSert=isUpSert)

        elif dataType == "Label":  # ?
            print("Save Label DataSeries")

        elif dataType == "Event":  # ?
            print("Save Event DataSeries")
        #
        print(str(datetime.now()), "Saved DataSeries:", dataSeriesName, "#Elements", str(dataSeries.Count()))


    def addItem(self, databaseName, collectionName, target, items):
        db = self.client[databaseName]
        coll = db[collectionName]
        update = {}
        update["$set"] = items  # {"Values.test3":3}
        coll.update_one(target, update, True)  # {"Key":"IC.CFE"}

    def addListItem(self, databaseName, collectionName, target, items):
        db = self.client[databaseName]
        coll = db[collectionName]
        update = {}
        update["$push"] = items  # {"Values":3}
        coll.update_one(target, update, True)  # {"Key":"IC.CFE"}

    def creatIndex(self, databaseName, collectionName, indexKeyName):
        db = self.client[databaseName]
        coll = db[collectionName]
        coll.create_index(indexKeyName)
        print("Created Index: " + indexKeyName + " for " + collectionName)

    def Close(self):
        self.client.close()

    def Delete(self, databaseName, collectionName, query):
        db = self.client[databaseName]
        coll = db[collectionName]
        coll.delete_many(query)
        # result = db.restaurants.delete_many({"borough": "Manhattan"})

    def DeleteCollections(self, databaseName, pattern = ""):
        db = self.client[databaseName]
        names = db.collection_names(False)
        for name in names:
            if pattern in name:
                coll = db[name]
                coll.drop()
                print("Delete Collection:" + name)

    def DeleteCollection(self, databaseName, collectionName):
        db = self.client[databaseName]
        names = db.collection_names(False)
        if collectionName in names:
            coll = db[collectionName]
            coll.drop()
            print("Delete Collection:" + collectionName)

    def deleteDataSeries(self, dataSeriesName):
        symbol = self.parseSymbol(dataSeriesName)
        type = self.parseDataType(dataSeriesName)
        name = self.parseDataName(dataSeriesName)

        if type == "Factor":
            self.delete("Factor", name, {"Symbol": symbol})

    def findWithFilter(self, databaseName, collectionName, filter={}, sort=None, projection={}):

        tempFilter = copy.deepcopy(filter)

        db = self.client[databaseName]
        #if self.password != None:
        #    db.authenticate(self.username, self.password)

        skip = 0
        limit = 0
        if "skip" in tempFilter:
            skip = tempFilter["skip"]
            tempFilter.pop("skip")
        if "limit" in tempFilter:
            limit = tempFilter["limit"]
            tempFilter.pop("limit")
        if "orderby" in tempFilter:
            sort = []
            for key, value in tempFilter["orderby"].items():
                if value == 1:
                    sortMethod = pymongo.ASCENDING
                else:
                    sortMethod = pymongo.DESCENDING
                sort.append((key, sortMethod))
            tempFilter.pop("orderby")

        #
        for key, value in tempFilter.items():
            if isinstance(value, dict):
                for sig, criteria in value.items():
                    if sig == ">":
                        tempFilter[key]["$gt"] = tempFilter[key][">"]
                        tempFilter[key].pop(">")
                    elif sig == ">=":
                        tempFilter[key]["$gte"] = tempFilter[key][">="]
                        tempFilter[key].pop(">=")
                    elif sig == "<":
                        tempFilter[key]["$lt"] = tempFilter[key]["<"]
                        tempFilter[key].pop("<")
                    elif sig == "<=":
                        tempFilter[key]["$lte"] = tempFilter[key]["<="]
                        tempFilter[key].pop("<=")

        #
        if len(projection) != 0:
            cursor = db[collectionName].find(tempFilter, skip=skip, limit=limit, sort=sort, projection=projection).max_scan(MAX_DATABASE_SCAN)
        else:
            cursor = db[collectionName].find(tempFilter, skip=skip, limit=limit, sort=sort).max_scan(MAX_DATABASE_SCAN)
        objList = []

        #
        for document in cursor:
            #print(document)
            jsonDoc = json.dumps(document, default=json_util.default)
            jsonObj = json.loads(jsonDoc, object_hook=json_util.object_hook)
            #if "_id" in jsonObj:
            #    del jsonObj["_id"]
            objList.append(jsonObj)
        return objList

    def find(self, databaseName, collectionName, beginDateTime=None, endDateTime=None, query={}, sort=None, asc=True, projection={}):
        #
        db = self.client[databaseName]
        query["StdDateTime"] = {}
        if beginDateTime != None:
            query["StdDateTime"]["$gte"] = beginDateTime
        if endDateTime != None:
            query["StdDateTime"]["$lte"] = endDateTime

        if not query["StdDateTime"]:
            del query["StdDateTime"]

        sortDoc = []
        if sort != None:
            sortDoc = sort
        else:
            if asc:
                sort = pymongo.ASCENDING
            else:
                sort = pymongo.DESCENDING
            sortDoc = [("StdDateTime", sort)]

        objList = self.findWithFilter(databaseName=databaseName,
                                      collectionName=collectionName,
                                      filter=query,
                                      sort=sortDoc,
                                      projection=projection)
        #
        return objList

    def findAll(self,databaseName, collectionName):
        db = self.client[databaseName]
        cursor = db[collectionName].find().max_scan(100000)
        objList = []
        for document in cursor:
            #print(document)
            jsonDoc = json.dumps(document, default=json_util.default)
            jsonObj = json.loads(jsonDoc, object_hook=json_util.object_hook)
            objList.append(jsonObj)
        return objList

    def getCollection(self, databaseName, collectionName):
        db = self.client[databaseName]
        coll =  db[collectionName]
        return coll

    def GetCollectionNames(self, databaseName, pattern=None):
        db = self.client[databaseName]
        names = db.collection_names(False)
        if pattern == None:
            return names
        #
        filteredNames = []
        for name in names:
            if pattern in name:
                filteredNames.append(name)
        return filteredNames

    def getDataSeries(self, dataSeriesName, datetime1=None, datetime2=None, instrumentType=None, asObject=True):

        #
        type = self.parseDataType(dataSeriesName)
        name = self.parseDataName(dataSeriesName)
        symbol = self.parseSymbol(dataSeriesName)
        #
        sort = [("StdDateTime", pymongo.ASCENDING)]
        query = {}

        #
        query["StdDateTime"] = {}
        if datetime1 != None:
            query["StdDateTime"]["$gte"] = datetime1
        if datetime2 != None:
            query["StdDateTime"]["$lte"] = datetime2
        if not query["StdDateTime"]:
            del query["StdDateTime"]

        # collectionName = name + "_Time_86400_Bar"
        if type == "Bar":
            collectionName = dataSeriesName
            # if instrumentType == None:
            #    # stockList = self.find("Instruments","Stock",fields={"Symbol":1})
            #    print("Unknown Bar Type")
            #    return
            if instrumentType == "Stock":  # and ( size == 86400 or size == 60):
                databaseName = "Quote"
            elif instrumentType == "Index":
                databaseName = "Index"
            elif instrumentType == "Future":
                databaseName = "Future"
                # query["DataSeries"] = dataSeriesName
            elif instrumentType == "Bond":
                databaseName = "Bond"
                # query["DataSeries"] = dataSeriesName
            elif instrumentType == "MutualFund":
                databaseName = "MutualFund"
            else:
                # ---default stock---
                databaseName = "Quote"
                # print("Unknown Bar Type")
                # return

        elif type == "Fundamental":
            databaseName = "Fundamental"
            collectionName = dataSeriesName

        elif type == "Portfolio":
            databaseName = "Portfolio"
            collectionName = "PortfolioHistData"

        elif type == "InstrumentList":
            databaseName = "Instruments"
            collectionName = "InstrumentList"
            query["Symbol"] = symbol

        elif type == "Factor":
            databaseName = "Factor"
            collectionName = name
            query["Symbol"] = symbol

        elif type == "Label":  # ?
            print("Load Label DataSeries")

        elif type == "Event":  # ?
            print("Load Event DataSeries")

        else:
            print("Unknown DataSerie Type: " + type)
            return

        #
        if asObject:
            # ---
            db = self.client[databaseName]
            # cursor = db[collectionName].find({"StdDateTime": {"$gt": beginDateTime, "$lt": endDateTime} }).max_scan(100000)
            cursor = db[collectionName].find(query).sort(sort)
            # ---
            dataSeries = DataSeries(dataSeriesName)
            for document in cursor:
                # print(document)
                jsonDoc = json.dumps(document, default=json_util.default)
                jsonObj = json.loads(jsonDoc, object_hook=json_util.object_hook)
                del jsonObj["_id"]
                dataSeries.Add(jsonObj)
        else:
            dataSeries = self.findWithFilter(databaseName, collectionName, query, sort)

        return dataSeries

    def increase(self, databaseName, collectionName, target, dataObject):
        db = self.client[databaseName]
        coll =  db[collectionName]
        update = {}
        update["$inc"] = dataObject
        coll.update_one(target, update, False)

    def Insert(self, databaseName, collectionName, dataObject):
        db = self.client[databaseName]
        coll = db[collectionName]
        coll.insert(dataObject)
        #coll.insert_one()

    def Insert_many(self, databaseName, collectionName, dataList):
        db = self.client[databaseName]
        coll = db[collectionName]
        coll.insert_many(dataList)

    def InsertIfNotExisted(self, databaseName, collectionName, target, dataObject):
        db = self.client[databaseName]
        coll =  db[collectionName]
        doc = coll.find_one(target)
        if doc == None:
            coll.insert(dataObject)
            return True
        return False

    def saveDataSeries(self, dataSeries, instrumentType=None, isUpSert=True):

        dataSeriesName = dataSeries.model_name
        print(str(datetime.now()) + " Saving DataSeries " + dataSeriesName)
        #
        if dataSeries.Count() == 0:
            print(dataSeriesName + " saveDataSeries::No Data to Save")
            return

        symbol = self.parseSymbol(dataSeriesName)
        type = self.parseDataType(dataSeriesName)
        name = self.parseDataName(dataSeriesName)

        databaseName = ""
        collectionName = ""
        keys = []
        insertType = 0
        if type == "Bar":
            size = self.parseSize(dataSeriesName)
            dataBaseStructureType = "BySymbol"
            if instrumentType == None:
                #stockList = self.find("Instruments","Stock",fields={"Symbol":1})
                print("Unknown Bar Type")
                return
            elif (instrumentType == "Stock") and (size == 86400 or size == 60 or size == 300):
                databaseName = "Quote"
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, databaseName, collectionName)
            elif instrumentType == "Index":
                databaseName = "Index"
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, databaseName, collectionName)
            elif instrumentType == "Future":
                databaseName = "Future"
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, databaseName, collectionName)
            elif instrumentType == "Bond":
                databaseName = "Bond"
                # query["DataSeries"] = dataSeriesName
            elif instrumentType == "MutualFund":
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, databaseName, collectionName)
            else:
                print("Unknown Bar Type")
                return

        elif type == "Fundamental":
            databaseName = "Fundamental"
            collectionName = dataSeriesName
            self.saveCollection(dataSeries.items, databaseName, collectionName)

        elif type == "InstrumentList":
            databaseName = "Instruments"
            collectionName = "InstrumentList"

        elif type == "Factor":
            databaseName = "Factor"
            collectionName = name
            self.saveCollection(dataSeries.items, databaseName, collectionName, keys=["Symbol"], isUpSert=isUpSert)

        elif type == "Label":# ?
            print("Save Label DataSeries")

        elif type == "Event":# ?
            print("Save Event DataSeries")
        #
        print(str(datetime.now()) + " Saved DataSeries " + dataSeriesName + " # of Elements " + str(dataSeries.Count()))

    def saveCollection(self, dataList, databaseName, collectionName, keys=[], isUpSert=True):

        # Loop to Insert DataBase
        # print(str(datetime.now()) + " Saving DataSeries " + collectionName)

        # ---Determine if the first time to save data (if the dataSeries exist?)---
        isFirstSaveData = False
        collectionNames = self.client[databaseName].collection_names(False)
        if collectionName not in collectionNames:
            isFirstSaveData = True
            isUpSert = False # First Save data, unnecessary to use Upsert Command

        # ---Create Key for each Element---
        for data in dataList:
            if "Key" not in data:
                localDateTime = Gadget.ToLocalDateTime(data["StdDateTime"])
                key = data["DataSeries"] + "_" + Gadget.ToDateTimeString(localDateTime) # Name + DateTime as identifier
                data["Key"] = key

        # if first to create DatSeries, Build Index
        if isFirstSaveData:
            db = self.client[databaseName]
            coll = db[collectionName]
            coll.create_index("Key")
            coll.create_index("StdDateTime")
            for key in keys:
                coll.create_index(key)
            #
            #self.insert_many(databaseName, collectionName, dataList)
            self.Upsert_Many(databaseName, collectionName, ["Key"], dataList)
        else:
            # for data in dataList:
            #    target = {"Key": data["Key"]}
            #    self.upsert(databaseName, collectionName, target, data)
            if isUpSert:
                self.Upsert_Many(databaseName, collectionName, ["Key"], dataList)
            else:
                self.Insert_Many(databaseName, collectionName, dataList)

        print(str(datetime.now()), "Saved Collection:", collectionName, "#Elements", str(len(dataList)))

    def Stats(self, databaseName, collectionName):
        db = self.client[databaseName]
        ret = db.command("collstats", collectionName)
        return ret

    def Upsert(self, databaseName, tableName, target, document):
        if "_id" in document:
            del document["_id"]
        #
        db = self.client[databaseName]
        coll = db[tableName]
        update = {}
        update["$set"] = document
        coll.update_one(target, update, True)

    def Upsert_Many(self, databaseName, collectionName, keys, dataList):
        db = self.client[databaseName]
        coll = db[collectionName]
        requests = []
        for data in dataList:
            filter = {}
            for key in keys:
                filter[key] = data[key]
            requests.append(pymongo.UpdateOne(filter, {"$set": data}, upsert=True))
        results = coll.bulk_write(requests)
        return results

    def Update(self, databaseName, collectionName, target, dataObject, delete=False):
        db = self.client[databaseName]
        coll = db[collectionName]
        update = {}
        if delete == True:
            update["$unset"] = dataObject
        else:
            update["$set"] = dataObject
        coll.update_one(target, update, False)


def CombineCollections(database, databaseName, fromPattern, toDatebaseName, toCollectionName):
    names = database.GetCollectionNames(databaseName, fromPattern)
    dataCount = 0
    count = 0
    skip = []
    skip.append("10000459.SH_Time_86400_Bar")
    skip.append("10000135.SH_Time_86400_Bar")
    for collectionName in names:
        count += 1
        symbol = collectionName.split("_")[0]
        # print("Process ", collectionName, count)
        #if count < 416:
        #    continue
        #if collectionName in skip:
        #    collections = database.Find(databaseName, collectionName)
        #    for data in collections:
        #        database.Upsert(databaseName, collectionName, {"Key":data["Key"]}, data)

        collections = database.Find(databaseName, collectionName)
        count1 = len(collections)
        dataCount += count1
        #database.saveCollection(collections, toDatebaseName, toCollectionName)

        collections2 = database.Find("Option", "DailyBars", {"Symbol": symbol})
        count2 = len(collections2)
        if count1 != count2:
            print("Process ", collectionName, count, count1, count2)
        pass
    #
    print("TotalBars", dataCount)


def CollectionTransit(fromDataBase, fromDataBaseName, fromCollectionName, toDataBase, toDatabaseName, toCollectionName, callBack=None):

    #
    fromCollection = fromDataBase.find(fromDataBaseName, fromCollectionName)

    #toDataBase.insert_many(toDatabaseName,toCollectionName,cursor)
    #
    count = 0
    for document in fromCollection:
        count += 1
        print(count)
        if callBack:
            newDocument = callBack(document)
        else:
            newDocument = document

        #
        if newDocument["Key"] != None:
            target = {"Key": newDocument["Key"]}
            toDataBase.upsert(toDatabaseName, toCollectionName, target, newDocument)
        else:
            toDataBase.insert(toDatabaseName, toCollectionName, newDocument)
        pass
    print("Transit Collection Done")


#---Copy Data To Another DB---
#---Filter, 只Copy特定collection_name名字的 collection
def DataBaseTransit(fromDataBase, toDataBase, databaseName, filter = None):

    db = fromDataBase.client[databaseName]
    db2 = toDataBase.client[databaseName]
    coll =  db.collection_names(False)
    process = False
    count = 0
    for collectionName in coll:
        if filter == None:
            process = True
        else:
            if filter in collectionName:
                process = True

        if process == True:
            cursor = db[collectionName].find()
            destCollection = db2[collectionName]
            #
            i = 0
            for document in cursor:
                #i = i + 1
                #print(i)
                del document["_id"]
                wcursor = destCollection.insert_one(document)
            #
            # coll.insert(cursor)#Inseert All
            # Create Index
            destCollection.create_index("Key")
            destCollection.create_index("StdDateTime")
            # CollectionTransit(fromDataBase, databaseName, collectionName, toDataBase, databaseName, collectionName)
            count = count + 1
            print(str(count) + " Copy DataSeries " + collectionName + " " + str(datetime.now()))


def RenameDataBase(database, oldName, newName):

    #db.copyDatabase(oldName, newName);
    db = database.client[oldName]
    coll = db.collection_names()
    #
    db = database.client["admin"]
    for collectionName in coll:
        commandStr = {}
        commandStr["renameCollection"] = oldName + "." + collectionName
        commandStr["to"] = newName + "." + collectionName
        db.command(commandStr)
        print("Renamed :" + collectionName)
        #db.runCommand({"renameCollection": oldName + ".000001.OF_Time_86400_Bar", "to": newName + ".000001.OF_Time_86400_Bar"});
    kkwood = 0


def CreateIndex(database, databaseName, collectionName, indexKeyName):
    database.creatIndex(databaseName, collectionName, indexKeyName)